Recover Binary Search Tree

Two elements of a binary search tree (BST) are swapped by mistake.

Recover the tree without changing its structure.

Note:

A solution using O(n) space is pretty straight forward. Could you devise a constant space solution?

Solution:

  1. /**
  2.  * Definition for a binary tree node.
  3.  * public class TreeNode {
  4.  *     int val;
  5.  *     TreeNode left;
  6.  *     TreeNode right;
  7.  *     TreeNode(int x) { val = x; }
  8.  * }
  9.  */
  10. public class Solution {
  11.   public void recoverTree(TreeNode root) {
  12.     if (root == null)
  13.       return;
  15.     TreeNode a = null;
  16.     TreeNode b = null;
  18.     TreeNode prev = null;
  19.     TreeNode curr = root;
  20.     Stack<TreeNode> stack = new Stack<>();
  22.     while (!stack.isEmpty() || curr != null) {
  23.       if (curr != null) {
  24.         stack.push(curr);
  25.         curr = curr.left;
  26.       } else {
  27.         curr = stack.pop();
  29.         if (prev != null && prev.val > curr.val) {
  30.           if (a == null) {
  31.             a = prev;
  32.           }
  34.           b = curr;
  35.         }
  37.         prev = curr;
  38.         curr = curr.right;
  39.       }
  40.     }
  42.     // swap a and b to fix the problem
  43.     if (a != null && b != null) {
  44.       int tmp = a.val;
  45.       a.val = b.val;
  46.       b.val = tmp;
  47.     }
  48.   }
  49. }